import os
import torch
import torch.nn as nn
from torchvision import transforms, models
from PIL import Image
import numpy as np

# 设备配置
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# ----------------------
# 1. 定义模型结构（必须与训练时完全一致）
# ----------------------
class ResNet50Binary(nn.Module):
    def __init__(self, pretrained=False):
        super(ResNet50Binary, self).__init__()
        self.model = models.resnet50(pretrained=False)
        num_features = self.model.fc.in_features
        self.model.fc = nn.Sequential(
            nn.Linear(num_features, 256),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(256, 1),
            nn.Sigmoid()
        )
        
    def forward(self, x):
        return self.model(x)

# ----------------------
# 2. 加载训练好的模型
# ----------------------
def load_resnet50_model(model_path):
    """加载训练好的ResNet50模型"""
    model = ResNet50Binary().to(device)
    model.load_state_dict(torch.load(model_path, map_location=device))
    model.eval()  # 切换到评估模式
    print(f"成功加载模型: {model_path}")
    return model

# ----------------------
# 3. 图像预处理（与训练时一致）
# ----------------------
def preprocess_image(image_path):
    """预处理单张图像"""
    transform = transforms.Compose([
        transforms.Resize((224, 224)),  # ResNet50的标准输入尺寸
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    img = Image.open(image_path).convert('RGB')
    return transform(img).unsqueeze(0).to(device)  # 增加batch维度

# ----------------------
# 4. 预测函数
# ----------------------
def predict_image(model, image_path, threshold=0.5):
    """
    使用ResNet50预测单张图像
    返回: {"label": "Real/Fake", "confidence": float, "is_real": bool}
    """
    # 预处理
    input_tensor = preprocess_image(image_path)
    
    # 预测
    with torch.no_grad():
        output = model(input_tensor)
        confidence = output.item()  # 直接获取sigmoid输出值
        is_real = confidence > threshold
    
    return {
        "label": "Real Face" if is_real else "Fake Face",
        "confidence": confidence,
        "is_real": is_real
    }

# ----------------------
# 5. 批量预测（可选）
# ----------------------
def predict_batch(model, image_dir, batch_size=32):
    """批量预测目录中的所有图像"""
    image_paths = [os.path.join(image_dir, f) for f in os.listdir(image_dir) 
                  if f.endswith(('.jpg', '.png', '.jpeg'))]
    
    results = []
    for img_path in image_paths:
        try:
            result = predict_image(model, img_path)
            result["image_path"] = img_path
            results.append(result)
        except Exception as e:
            print(f"预测失败: {img_path} - {str(e)}")
    
    return results

# ----------------------
# 6. 使用示例
# ----------------------
if __name__ == "__main__":
    # 加载模型（替换为你的模型路径）
    MODEL_PATH = "best_model.pth"  # 或 'checkpoints/resnet50_final.pth'
    model = load_resnet50_model(MODEL_PATH)
    
    # 单张图像预测
    test_image = "test.jpg"  # 替换为你的测试图像路径
    result = predict_image(model, test_image)
    print(f"\n预测结果: {result['label']} (置信度: {result['confidence']:.4f})")
    
    # 批量预测（可选）
    # batch_results = predict_batch(model, "test_images/")
    # for res in batch_results:
    #     print(f"{res['image_path']}: {res['label']} ({res['confidence']:.2f})")
